package pers.jsan.litheopen.security.config.request;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.CharArrayWriter;
import java.io.IOException;
import java.io.PrintWriter;

/**
 * @author 3hyzy
 * @data 2020-04-16
 **/
public class BodyReaderHttpServletResponseWrapper extends HttpServletResponseWrapper {


    private PrintWriter cachedWriter;
    private CharArrayWriter bufferedWriter;

    private ByteArrayOutputStream buffer;
    private ServletOutputStream out;

    public BodyReaderHttpServletResponseWrapper(HttpServletResponse response) {
        super(response);

        buffer = new ByteArrayOutputStream();
        out = new WrapperOutputStream(buffer);

        bufferedWriter = new CharArrayWriter();
        cachedWriter = new PrintWriter(bufferedWriter);
    }

    @Override
    public void flushBuffer() throws IOException {
        if (out != null) {
            out.flush();
        }
        if(buffer!=null){
            buffer.flush();
        }
    }


    @Override
    public PrintWriter getWriter() {
        return cachedWriter;
    }

    public String getResult() {
        try {
            flushBuffer();
        } catch (IOException e) {
            e.printStackTrace();
        }
        if(buffer.size()>0){
            return new String(buffer.toByteArray());
        }else {
            return bufferedWriter.toString();
        }

    }

    class WrapperOutputStream extends ServletOutputStream {
        private ByteArrayOutputStream bos;

        public WrapperOutputStream(ByteArrayOutputStream bos) {
            this.bos = bos;
        }

        @Override
        public void write(int b) throws IOException {
            bos.write(b);
        }

        @Override
        public boolean isReady() {
            return false;
        }

        @Override
        public void setWriteListener(WriteListener arg0) {

        }
    }
}
